import logging
import numpy as np
import pandas as pd
from qtorch.core.qstrategy import QStrategy

class MultiIndicatorStrategy(QStrategy):
    """
    多指标组合策略，结合多种技术指标进行交易决策
    
    指标包括：
    1. 移动平均线（短期和长期）
    2. RSI（相对强弱指标）
    3. MACD（移动平均收敛/发散）
    4. 布林带
    
    参数:
        fast_ma: 快速移动平均线周期
        slow_ma: 慢速移动平均线周期
        rsi_period: RSI计算周期
        rsi_overbought: RSI超买阈值
        rsi_oversold: RSI超卖阈值
        macd_fast: MACD快线周期
        macd_slow: MACD慢线周期
        macd_signal: MACD信号线周期
        bb_period: 布林带周期
        bb_std: 布林带标准差倍数
        weight_ma: 移动平均线权重
        weight_rsi: RSI权重
        weight_macd: MACD权重
        weight_bb: 布林带权重
        threshold: 组合信号阈值，大于此值买入，小于负值卖出
    """
    
    def __init__(self, fast_ma=20, slow_ma=50, 
                rsi_period=14, rsi_overbought=70, rsi_oversold=30,
                macd_fast=12, macd_slow=26, macd_signal=9,
                bb_period=20, bb_std=2,
                weight_ma=1.0, weight_rsi=1.0, weight_macd=1.0, weight_bb=1.0,
                threshold=0.5):
        """初始化多指标策略"""
        super().__init__()
        # 移动平均线参数
        self.fast_ma = fast_ma
        self.slow_ma = slow_ma
        
        # RSI参数
        self.rsi_period = rsi_period
        self.rsi_overbought = rsi_overbought
        self.rsi_oversold = rsi_oversold
        
        # MACD参数
        self.macd_fast = macd_fast
        self.macd_slow = macd_slow
        self.macd_signal = macd_signal
        
        # 布林带参数
        self.bb_period = bb_period
        self.bb_std = bb_std
        
        # 指标权重
        self.weight_ma = weight_ma
        self.weight_rsi = weight_rsi
        self.weight_macd = weight_macd
        self.weight_bb = weight_bb
        
        # 组合信号阈值
        self.threshold = threshold
        
        # 缓存计算结果
        self._bb_upper = None
        self._bb_lower = None
        self._bb_middle = None
        self._prev_macd_hist = 0
        
    def initialize(self):
        """策略初始化"""
        logging.info("初始化多指标组合策略")
        logging.info(f"移动平均线: 快线={self.fast_ma}, 慢线={self.slow_ma}, 权重={self.weight_ma}")
        logging.info(f"RSI: 周期={self.rsi_period}, 超买={self.rsi_overbought}, 超卖={self.rsi_oversold}, 权重={self.weight_rsi}")
        logging.info(f"MACD: 快线={self.macd_fast}, 慢线={self.macd_slow}, 信号线={self.macd_signal}, 权重={self.weight_macd}")
        logging.info(f"布林带: 周期={self.bb_period}, 标准差={self.bb_std}, 权重={self.weight_bb}")
        logging.info(f"信号阈值: {self.threshold}")
        
    def next(self):
        """
        处理当前数据点，生成交易信号
        
        返回:
            int: 1(买入)、0(持有)、-1(卖出)
        """
        # 检查数据是否充足
        if len(self.prices) < max(self.slow_ma, self.rsi_period, self.macd_slow + self.macd_signal, self.bb_period):
            return 0  # 数据不足
            
        # 计算各个指标信号
        ma_signal = self._get_ma_signal()
        rsi_signal = self._get_rsi_signal()
        macd_signal = self._get_macd_signal()
        bb_signal = self._get_bb_signal()
        
        # 如果有指标计算失败，跳过
        if None in [ma_signal, rsi_signal, macd_signal, bb_signal]:
            return 0
            
        # 计算加权组合信号（范围：-1到1）
        total_weight = self.weight_ma + self.weight_rsi + self.weight_macd + self.weight_bb
        if total_weight == 0:
            return 0  # 避免除以零
            
        combo_signal = (
            ma_signal * self.weight_ma + 
            rsi_signal * self.weight_rsi + 
            macd_signal * self.weight_macd + 
            bb_signal * self.weight_bb
        ) / total_weight
        
        # 根据组合信号和阈值生成交易信号
        if combo_signal > self.threshold:
            return 1  # 买入信号
        elif combo_signal < -self.threshold:
            return -1  # 卖出信号
        else:
            return 0  # 持有信号
    
    def _get_ma_signal(self):
        """计算移动平均线信号"""
        try:
            fast_ma = self.SMA(self.fast_ma)
            slow_ma = self.SMA(self.slow_ma)
            
            if fast_ma is None or slow_ma is None:
                return None
                
            # 快线在慢线上方为正信号，下方为负信号
            if fast_ma > slow_ma:
                return 1.0
            elif fast_ma < slow_ma:
                return -1.0
            else:
                return 0.0
        except Exception as e:
            logging.warning(f"计算MA信号错误: {str(e)}")
            return None
    
    def _get_rsi_signal(self):
        """计算RSI信号"""
        try:
            rsi = self.RSI(self.rsi_period)
            
            if rsi is None:
                return None
                
            # RSI低于超卖线为买入信号，高于超买线为卖出信号
            # 将RSI值归一化到[-1, 1]的范围
            if rsi <= self.rsi_oversold:
                # 超卖区域，买入信号
                normalized = -1.0 + 2.0 * (rsi / self.rsi_oversold)
                return -normalized  # 反转信号使超卖接近1.0
            elif rsi >= self.rsi_overbought:
                # 超买区域，卖出信号
                normalized = -1.0 + 2.0 * ((100 - rsi) / (100 - self.rsi_overbought))
                return -normalized  # 反转信号使超买接近-1.0
            else:
                # 中性区域，线性映射
                normalized = -1.0 + 2.0 * ((rsi - self.rsi_oversold) / (self.rsi_overbought - self.rsi_oversold))
                return -normalized  # 反转信号
        except Exception as e:
            logging.warning(f"计算RSI信号错误: {str(e)}")
            return None
    
    def _get_macd_signal(self):
        """计算MACD信号"""
        try:
            macd_result = self.MACD(self.macd_fast, self.macd_slow, self.macd_signal)
            
            if macd_result is None:
                return None
                
            macd_line = macd_result['macd']
            signal_line = macd_result['signal']
            histogram = macd_result['histogram']
            
            # MACD交叉信号
            if macd_line > signal_line:
                cross_signal = 1.0
            elif macd_line < signal_line:
                cross_signal = -1.0
            else:
                cross_signal = 0.0
                
            # 柱状图变化信号（趋势加速/减速）
            hist_change = 0.0
            if self._prev_macd_hist != 0:
                if histogram > self._prev_macd_hist:
                    hist_change = 0.5  # 上升
                elif histogram < self._prev_macd_hist:
                    hist_change = -0.5  # 下降
            self._prev_macd_hist = histogram
            
            # 组合MACD信号
            return 0.7 * cross_signal + 0.3 * hist_change
            
        except Exception as e:
            logging.warning(f"计算MACD信号错误: {str(e)}")
            return None
    
    def _get_bb_signal(self):
        """计算布林带信号"""
        try:
            # 计算布林带
            middle = self.SMA(self.bb_period)
            if middle is None or len(self.prices) < self.bb_period:
                return None
                
            # 计算标准差
            price_array = np.array(self.prices[-self.bb_period:])
            std = np.std(price_array)
            
            # 布林带上下轨
            upper = middle + self.bb_std * std
            lower = middle - self.bb_std * std
            
            # 保存计算结果
            self._bb_upper = upper
            self._bb_lower = lower
            self._bb_middle = middle
            
            # 当前价格
            current_price = self.prices[-1]
            
            # 计算带宽百分比位置
            band_width = upper - lower
            if band_width == 0:
                return 0.0
                
            # 价格在带中的位置百分比（0表示在下轨，1表示在上轨）
            position = (current_price - lower) / band_width
            
            # 将位置映射到[-1, 1]范围，越接近下轨越接近1（买入信号），越接近上轨越接近-1（卖出信号）
            return 1.0 - 2.0 * position
            
        except Exception as e:
            logging.warning(f"计算布林带信号错误: {str(e)}")
            return None
    
    def add_indicator_data(self, data):
        """
        向数据中添加策略使用的指标，用于可视化
        
        参数:
            data: DataFrame，回测数据
            
        返回:
            DataFrame: 添加了指标列的数据
        """
        # 添加移动平均线
        data[f'SMA_{self.fast_ma}'] = data['close'].rolling(window=self.fast_ma).mean()
        data[f'SMA_{self.slow_ma}'] = data['close'].rolling(window=self.slow_ma).mean()
        
        # 添加RSI
        delta = data['close'].diff()
        gain = delta.where(delta > 0, 0)
        loss = -delta.where(delta < 0, 0)
        avg_gain = gain.rolling(window=self.rsi_period).mean()
        avg_loss = loss.rolling(window=self.rsi_period).mean()
        rs = avg_gain / avg_loss
        data['RSI'] = 100 - (100 / (1 + rs))
        
        # 添加MACD
        ema_fast = data['close'].ewm(span=self.macd_fast, adjust=False).mean()
        ema_slow = data['close'].ewm(span=self.macd_slow, adjust=False).mean()
        data['MACD'] = ema_fast - ema_slow
        data['MACD_Signal'] = data['MACD'].ewm(span=self.macd_signal, adjust=False).mean()
        data['MACD_Hist'] = data['MACD'] - data['MACD_Signal']
        
        # 添加布林带
        data['BB_Middle'] = data['close'].rolling(window=self.bb_period).mean()
        data['BB_Std'] = data['close'].rolling(window=self.bb_period).std()
        data['BB_Upper'] = data['BB_Middle'] + self.bb_std * data['BB_Std']
        data['BB_Lower'] = data['BB_Middle'] - self.bb_std * data['BB_Std']
        
        return data